from PyQt6.QtWidgets import QPlainTextEdit, QWidget, QTextEdit
from PyQt6.QtGui import (QTextCharFormat, QSyntaxHighlighter, QPainter, 
                        QColor, QTextFormat)
from PyQt6.QtCore import Qt, QRect, QSize
import os
from pygments import highlight
from pygments.lexers import get_lexer_for_filename, TextLexer
from pygments.formatters import HtmlFormatter
from pygments.styles import get_style_by_name
from pygments.util import ClassNotFound

class LineNumberArea(QWidget):
    def __init__(self, editor):
        super().__init__(editor)
        self.editor = editor

    def sizeHint(self):
        return QSize(self.editor.line_number_area_width(), 0)

    def paintEvent(self, event):
        self.editor.lineNumberAreaPaintEvent(event)

class CodeViewer(QPlainTextEdit):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setReadOnly(True)
        self.setLineWrapMode(QPlainTextEdit.LineWrapMode.NoWrap)
        
        # 设置字体
        self.setFont('Consolas', 11)
        
        # 创建行号区域
        self.line_number_area = LineNumberArea(self)
        
        # 连接信号
        self.blockCountChanged.connect(self.update_line_number_area_width)
        self.updateRequest.connect(self.update_line_number_area)
        self.cursorPositionChanged.connect(self.highlight_current_line)
        
        # 初始化行号区域宽度
        self.update_line_number_area_width(0)
        
        # 高亮当前行
        self.highlight_current_line()

    def setFont(self, family, size):
        """设置编辑器字体"""
        from PyQt6.QtGui import QFont
        font = QFont(family, size)
        font.setFixedPitch(True)
        super().setFont(font)

    def line_number_area_width(self):
        """计算行号区域的宽度"""
        digits = 1
        max_num = max(1, self.blockCount())
        while max_num >= 10:
            max_num //= 10
            digits += 1
        space = 3 + self.fontMetrics().horizontalAdvance('9') * digits
        return space

    def update_line_number_area_width(self, newBlockCount):
        """更新行号区域宽度"""
        self.setViewportMargins(self.line_number_area_width(), 0, 0, 0)

    def update_line_number_area(self, rect, dy):
        """更新行号区域"""
        if dy:
            self.line_number_area.scroll(0, dy)
        else:
            self.line_number_area.update(0, rect.y(), 
                self.line_number_area.width(), rect.height())
        if rect.contains(self.viewport().rect()):
            self.update_line_number_area_width(0)

    def resizeEvent(self, event):
        """重写改变大小事件"""
        super().resizeEvent(event)
        cr = self.contentsRect()
        self.line_number_area.setGeometry(QRect(cr.left(), cr.top(),
            self.line_number_area_width(), cr.height()))

    def lineNumberAreaPaintEvent(self, event):
        """绘制行号区域"""
        painter = QPainter(self.line_number_area)
        painter.fillRect(event.rect(), Qt.GlobalColor.lightGray)
        
        block = self.firstVisibleBlock()
        blockNumber = block.blockNumber()
        top = round(self.blockBoundingGeometry(block).translated(
            self.contentOffset()).top())
        bottom = top + round(self.blockBoundingRect(block).height())

        while block.isValid() and top <= event.rect().bottom():
            if block.isVisible() and bottom >= event.rect().top():
                number = str(blockNumber + 1)
                painter.setPen(Qt.GlobalColor.black)
                painter.drawText(0, top, self.line_number_area.width(),
                    self.fontMetrics().height(),
                    Qt.AlignmentFlag.AlignRight, number)
            block = block.next()
            top = bottom
            bottom = top + round(self.blockBoundingRect(block).height())
            blockNumber += 1

    def highlight_current_line(self):
        """高亮当前行"""
        extraSelections = []
        if not self.isReadOnly():
            selection = QTextEdit.ExtraSelection()
            lineColor = QColor(Qt.GlobalColor.yellow).lighter(160)
            selection.format.setBackground(lineColor)
            selection.format.setProperty(QTextFormat.Property.FullWidthSelection, True)
            selection.cursor = self.textCursor()
            selection.cursor.clearSelection()
            extraSelections.append(selection)
        self.setExtraSelections(extraSelections)

    def load_file(self, file_path):
        """加载并显示文件内容"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            
            # 根据文件扩展名选择合适的语法高亮器
            try:
                lexer = get_lexer_for_filename(file_path)
            except ClassNotFound:
                lexer = TextLexer()
            
            # 使用Pygments进行语法高亮
            formatter = HtmlFormatter(style='monokai')
            highlighted_content = highlight(content, lexer, formatter)
            
            # 设置高亮后的内容
            self.setHtml(highlighted_content)
            return True
        except Exception as e:
            print(f"Error loading file: {str(e)}")
            return False

class TabViewer:
    """管理代码查看器的标签页"""
    def __init__(self, tab_widget):
        self.tab_widget = tab_widget
        self.opened_files = {}  # 存储已打开的文件 {file_path: tab_index}
        
    def open_file(self, file_path):
        """在新标签页中打开文件"""
        # 如果文件已经打开，切换到对应标签
        if file_path in self.opened_files:
            self.tab_widget.setCurrentIndex(self.opened_files[file_path])
            return
            
        # 创建新的代码查看器
        viewer = CodeViewer()
        if viewer.load_file(file_path):
            # 添加新标签
            tab_index = self.tab_widget.addTab(viewer, os.path.basename(file_path))
            self.opened_files[file_path] = tab_index
            self.tab_widget.setCurrentIndex(tab_index)
            
            # 添加关闭按钮的响应
            self.tab_widget.setTabsClosable(True)
            self.tab_widget.tabCloseRequested.connect(self.close_tab)
    
    def close_tab(self, index):
        """关闭标签页"""
        # 从已打开文件字典中移除
        file_path = next((path for path, idx in self.opened_files.items() 
                         if idx == index), None)
        if file_path:
            del self.opened_files[file_path]
            
        # 更新其他标签的索引
        for path, idx in self.opened_files.items():
            if idx > index:
                self.opened_files[path] = idx - 1
                
        # 移除标签页
        self.tab_widget.removeTab(index)
